from typing import Any, Callable, Iterable, TypeVar

from zkl_aiutils_training import Dock, DockPlugin

from zkl_ptutils_training.plugins import MetricOrValue, ValidationProvider
from zkl_ptutils_training.tasks import MLProcessor, MLProducer
from .engine import DelegatingInferringEngine, MLInferringEngine, MLTrainingEngine
from .processing import MLValidatingTask

AnyInput = TypeVar('AnyInput')
AnyOutput = TypeVar('AnyOutput')
AnyCreateArgs = TypeVar('AnyCreateArgs')


class StandardTrainer(MLProcessor[AnyInput, AnyOutput], DockPlugin):
    def __init__(self, engine: MLTrainingEngine[AnyInput, AnyOutput]):
        self._engine = engine

    def on_installed(self, dock: Dock):
        super().on_installed(dock)
        dock.install(self._engine)

    def process(self, input: AnyInput) -> AnyOutput:
        output = self.forward(input)
        self.backward(output)
        self.optimize()
        return output

    def forward(self, input: AnyInput) -> AnyOutput:
        return self._engine.forward(input)

    def backward(self, output: AnyOutput):
        return self._engine.backward(output)

    def optimize(self) -> AnyOutput:
        return self._engine.optimize()

    def validator(self) -> 'StandardValidator[AnyInput, AnyOutput]':
        engine = DelegatingInferringEngine(self._engine)
        return StandardValidator(engine)


class StandardValidator(MLProcessor[AnyInput, AnyOutput], DockPlugin):
    def __init__(self, engine: MLInferringEngine[AnyInput, AnyOutput]):
        self._engine = engine

    def on_installed(self, dock: Dock):
        super().on_installed(dock)
        dock.install(self._engine)

    def process(self, input: AnyInput) -> AnyOutput:
        return self.forward(input)

    def forward(self, input: AnyInput) -> AnyOutput:
        return self._engine.forward(input)


class StandardValidationProvider(ValidationProvider, DockPlugin):
    def __init__(self,
        producer_factory: Callable[[], MLProducer[AnyInput]],
        plugins_factory: Callable[[], Iterable[Any]] | None = None
    ):
        super().__init__()
        self._producer_factory = producer_factory
        self._plugins_factory = plugins_factory

    @property
    def trainer(self) -> StandardTrainer:
        return self.dock.get_plugin(StandardTrainer)

    def validate(self) -> dict[str, MetricOrValue]:
        dock = Dock()

        task = MLValidatingTask(
            producer=self._producer_factory(),
            processor=self.trainer.validator())
        dock.install(task)

        if self._plugins_factory:
            for plugin in self._plugins_factory():
                dock.install(plugin)

        return task.run()
